#!/usr/bin/env python
#
# testing for keep-alive connection
#   agent-pool support
#

import time,socket,base64,fcntl,errno,sys
import threading
import Queue as queue

import celtuce.net.ioloop as ioloop
import celtuce.net.iostream as iostream
import celtuce.help.codec as codec

from celtuce.gatepacket import Packet
from celtuce.help import logging


#######################################################################

class   Client(object) :
    def __init__(self, server_host, server_port) :
        self._server_host = server_host
        self._server_port = server_port

        self._END_TAG = "\r\n\r\n"

        self._socket = None
        self._stream = None

        self._keep_alive = True

        self._logger = logging.getLogger('client')
        ## __init__()

    def __del__(self) :
        self.__disconnect(force=True)
        ## __del__()

    ####-----------------------------------------------------------####

    def __connect(self) :
        if self._socket : 
            #self._stream = iostream.IOStream( self._socket, self._io_loop )
            return

        try :
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
            s.connect( (self._server_host, self._server_port) )
        except :
            self._logger.excpt()
            raise
        self._socket = s

        self._io_loop = ioloop.IOLoop()
        self._stream = iostream.IOStream( self._socket, self._io_loop )
        self._stream.set_close_callback( self.onDisconnect )

        return
        ## __connect()

    def __disconnect(self, force=False) :
        if self._keep_alive and not force :
            return

        if self._stream :
            self._stream.close()
            self._stream = None

        if self._socket :
            try :
                self._socket.close()
            except :
                self._logger.excpt()

        self._socket = None
        self._io_loop = None
        ## __disconnect()

    def onDisconnect(self) :
        self._logger.info('socket closed.')
        self._io_loop.stop()
        self.__disconnect(force=True)
        ## onDisconnect()

    ####-----------------------------------------------------------####

    def rpc(self, request) :
        self.__connect()

        #self._logger.debug('write request : %s', request)
        self._stream.write( ''.join((request, self._END_TAG)), self.__onWriteComplete )
        self._stream.read_until(self._END_TAG, self.__onResponse)
        self._response = None

        self._io_loop.start()

        self.__disconnect()
        return self._response
        ## rpc()

    def __onWriteComplete(self) :
        #self._logger.debug('write done.')
        pass
        ## __onWriteComplete()

    def __onResponse(self, response_str) :
        #self._logger.debug('read response : %s', repr(response_str))
        response = response_str.strip()
        if response :
            self._response = response_str.strip()
            self._io_loop.stop()
        else :
            self._stream.read_until(self._END_TAG, self.__onResponse)
        ## __onResponse()

    ## class Client


#######################################################################


def echo(any_data) :
    return any_data.strip()
    ## echo()


TCP_PORT = 14390

TEST_TIMES = 10000

def runTest() :
    logger = logging.getLogger('test')

    client = Client('', TCP_PORT)

    try :
        # ping
        packet = Packet()
        packet.action = Packet.ACTION_REQUEST
        packet.table = 'test.gate4'
        packet.op = Packet.OP_INSERT
        packet.key = 'hello'
        packet.data = 'world'
        request,request_id = Packet.enpack( packet )
        while True :
            response_str = client.rpc( request )
            response = Packet.depack( response_str )
            if response.isOk() :
                break
            else :
                logger.error('ERROR : %s, and retry ...', response.error)
                time.sleep(1.0)

        logger.info('REQUEST  : %s', request)
        logger.info('RESPONSE : %s', response)

        # test insert 
        key_prefix = 'hello'
        data_prefix = 'world'
        time_start = time.time()
        times = TEST_TIMES
        sum_bytes = 0
        for i in range(times) :
            packet = Packet()
            packet.action = Packet.ACTION_REQUEST
            packet.table = 'test.gate4'
            packet.op = Packet.OP_INSERT
            packet.key = '-'.join((key_prefix, str(i)))
            packet.data = '-'.join((data_prefix, str(i)))
            request,request_id = Packet.enpack( packet )
            response_str = client.rpc( request )
            response = Packet.depack( response_str )
            if not response.isOk() :
                print response
                raise RuntimeError

            sum_bytes += len(response_str)
        time_last = time.time() - time_start

        tps = times / time_last
        bw = sum_bytes / time_last / 1024.0
        logger.info('rpc-INSERT %d times last %.2f seconds, tps:%.2f, bw:%.2f KB/s',times,time_last,tps,bw)

        # test find
        key_prefix = 'hello'
        data_prefix = 'world'
        time_start = time.time()
        times = TEST_TIMES
        sum_bytes = 0
        count_success = 0
        count_failure = 0
        for i in range(times) :
            packet = Packet()
            packet.action = Packet.ACTION_REQUEST
            packet.table = 'test.gate4'
            packet.op = Packet.OP_FIND
            packet.key = '-'.join((key_prefix, str(i)))
            data_want = '-'.join((data_prefix, str(i)))
            request,request_id = Packet.enpack( packet )
            response_str = client.rpc( request )
            response = Packet.depack( response_str )
            if data_want == response.data :
                count_success += 1
            else :
                print 'GOT %r, but WANT %r' % (response.data, data_want)
                count_failure += 1

            sum_bytes += len(response_str)
        time_last = time.time() - time_start

        tps = times / time_last
        bw = sum_bytes / time_last / 1024.0
        logger.info('rpc-FIND %d times last %.2f seconds, tps:%.2f, bw:%.2f KB/s count(success/failure)=%d/%d',times,time_last,tps,bw,count_success,count_failure)

        # test erase 
        key_prefix = 'hello'
        data_prefix = 'world'
        time_start = time.time()
        times = TEST_TIMES
        sum_bytes = 0
        for i in range(times) :
            packet = Packet()
            packet.action = Packet.ACTION_REQUEST
            packet.table = 'test.gate4'
            packet.op = Packet.OP_ERASE
            packet.key = '-'.join((key_prefix, str(i)))
            packet.data = '-'.join((data_prefix, str(i)))
            request,request_id = Packet.enpack( packet )
            response_str = client.rpc( request )
            response = Packet.depack( response_str )

            sum_bytes += len(response_str)
        time_last = time.time() - time_start

        tps = times / time_last
        bw = sum_bytes / time_last / 1024.0
        logger.info('rpc-ERASE %d times last %.2f seconds, tps:%.2f, bw:%.2f KB/s',times,time_last,tps,bw)

        # test find
        key_prefix = 'hello'
        data_prefix = 'world'
        time_start = time.time()
        times = TEST_TIMES
        sum_bytes = 0
        count_success = 0
        count_failure = 0
        for i in range(times) :
            packet = Packet()
            packet.action = Packet.ACTION_REQUEST
            packet.table = 'test.gate4'
            packet.op = Packet.OP_FIND
            packet.key = '-'.join((key_prefix, str(i)))
            data_want = None
            request,request_id = Packet.enpack( packet )
            response_str = client.rpc( request )
            response = Packet.depack( response_str )
            if data_want == response.data :
                count_success += 1
            else :
                count_failure += 1

            sum_bytes += len(response_str)
        time_last = time.time() - time_start

        tps = times / time_last
        bw = sum_bytes / time_last / 1024.0
        logger.info('rpc-FIND %d times last %.2f seconds, tps:%.2f, bw:%.2f KB/s count(success/failure)=%d/%d',times,time_last,tps,bw,count_success,count_failure)

    finally :
        logger.info('shutdown ...')

    ## runTest()


########################################################################

if  __name__ == '__main__' :

    runTest()


